{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b9f48e9b",
   "metadata": {},
   "source": [
    "# JIT Compilation with `BrainPyObject`\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_advanced/compilation.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_advanced/compilation.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355bb9b6",
   "metadata": {},
   "source": [
    "@[Chaoming Wang](https://github.com/chaoming0625)\n",
    "@[Xiaoyu Chen](mailto:c-xy17@tsinghua.org.cn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a625b0ab",
   "metadata": {},
   "source": [
    "In this section, we are going to talk about code compilation that can accelerate your model running performance. "
   ]
  },
  {
   "cell_type": "code",
   "id": "13e791f8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:31:58.780972Z",
     "start_time": "2025-10-06T03:31:58.777546Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "# bm.set_platform('cpu')"
   ],
   "outputs": [],
   "execution_count": 28
  },
  {
   "cell_type": "markdown",
   "id": "cdfd2be7",
   "metadata": {},
   "source": [
    "## ``brainpy.math.jit()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "123027f3",
   "metadata": {},
   "source": [
    "[JAX](https://github.com/google/jax) provides JIT compilation ``jax.jit()`` for [pure functions](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions).In most cases, however, we code with Python classes. ``brainpy.math.jit()`` is intended to extend just-in-time compilation to class objects. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9406eacd",
   "metadata": {},
   "source": [
    "### JIT compilation for class objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fae4adb",
   "metadata": {},
   "source": [
    "The constraints for class-object JIT ciompilation include:\n",
    "\n",
    "- The JIT target must be a subclass of ``brainpy.BrainPyObject``.\n",
    "- Dynamically changed variables must be labeled as ``brainpy.math.Variable``.\n",
    "- Updating Variables must be accomplished by in-place operations. "
   ]
  },
  {
   "cell_type": "code",
   "id": "a5374857",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:31:58.800937Z",
     "start_time": "2025-10-06T03:31:58.796858Z"
    }
   },
   "source": [
    "class LogisticRegression(bp.BrainPyObject):\n",
    "    def __init__(self, dimension):\n",
    "        super(LogisticRegression, self).__init__()\n",
    "\n",
    "        # parameters\n",
    "        self.dimension = dimension\n",
    "\n",
    "        # variables\n",
    "        self.w = bm.Variable(2.0 * bm.ones(dimension) - 1.3)\n",
    "\n",
    "    def __call__(self, X, Y):\n",
    "        u = bm.dot(((1.0 / (1.0 + bm.exp(-Y * bm.dot(X, self.w))) - 1.0) * Y), X)\n",
    "        self.w[:] = self.w - u \n",
    "        # The above line can also be expressed as: \n",
    "        # \n",
    "        #   self.w.value = self.w - u \n",
    "        # \n",
    "        # or, \n",
    "        # \n",
    "        #   self.w.update(self.w - u)"
   ],
   "outputs": [],
   "execution_count": 29
  },
  {
   "cell_type": "markdown",
   "id": "2440bf0e",
   "metadata": {},
   "source": [
    "In this example, weight *self.w* is a dynamically changed variable, thus marked as ``Variable``. During the update phase ``__call__()``, *self.w* is in-place updated through ``self.w[:] = ...``. Alternatively, one can replace the data in the variable by ``self.w.value = ...`` or ``self.w.update(...)``."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "893ec359",
   "metadata": {},
   "source": [
    "Now this logistic regression can be accelerated by JIT compilation."
   ]
  },
  {
   "cell_type": "code",
   "id": "462f745b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:31:58.814578Z",
     "start_time": "2025-10-06T03:31:58.805947Z"
    }
   },
   "source": [
    "num_dim, num_points = 10, 200000\n",
    "points = bm.random.random((num_points, num_dim))\n",
    "labels = bm.random.random(num_points)"
   ],
   "outputs": [],
   "execution_count": 30
  },
  {
   "cell_type": "code",
   "id": "f5e5b98c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:09.549674Z",
     "start_time": "2025-10-06T03:31:58.826599Z"
    }
   },
   "source": [
    "lr = LogisticRegression(10)\n",
    "\n",
    "%timeit lr(points, labels)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.32 ms ± 24.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "execution_count": 31
  },
  {
   "cell_type": "code",
   "id": "c70b1eba",
   "metadata": {
    "scrolled": true,
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.398725Z",
     "start_time": "2025-10-06T03:32:09.584398Z"
    }
   },
   "source": [
    "lr_jit = bm.jit(lr)\n",
    "\n",
    "%timeit lr_jit(points, labels)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.33 ms ± 11.1 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "execution_count": 32
  },
  {
   "cell_type": "markdown",
   "id": "25a3b576",
   "metadata": {},
   "source": [
    "### JIT mechanism"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1046bd0",
   "metadata": {},
   "source": [
    "The mechanism of JIT compilation is that BrainPy automatically transforms your class methods into functions. \n",
    "\n",
    "``brainpy.math.jit()`` receives a ``dyn_vars`` argument, which denotes the dynamically changed variables. If it is not provided by users, BrainPy will automatically detect them by calling ``BrainPyObject.vars()`` (only variables labeled as `Variable` will be automatically retrieved by BrainPyObject.vars()). Once receiving \"dyn_vars\", BrainPy will treat \"dyn_vars\" as function arguments and then transform class objects into functions."
   ]
  },
  {
   "cell_type": "code",
   "id": "42fbe267",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.446878Z",
     "start_time": "2025-10-06T03:32:20.440516Z"
    }
   },
   "source": [
    "lr_jit"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.warp_to_no_state_input_output.<locals>.wrapper(X, Y)>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 33
  },
  {
   "cell_type": "markdown",
   "id": "e65f0e22",
   "metadata": {},
   "source": [
    "Therefore, the secrete of ``brainpy.math.jit()`` is providing \"dyn_vars\". No matter your target is a class object, a method in the class object, or a pure function, if there are dynamically changed variables, you just pack them into ``brainpy.math.jit()`` as \"dyn_vars\". Then, all the compilation and acceleration will be handled by BrainPy automatically. Let's illustrate this by several examples. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29d5d84",
   "metadata": {},
   "source": [
    "#### Example 1: JIT compiled methods in a class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02e79a7f",
   "metadata": {},
   "source": [
    "In this example, we try to run a method just-in-time in a class, in which the object variable are used to compute the final results. "
   ]
  },
  {
   "cell_type": "code",
   "id": "076fc88b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.467212Z",
     "start_time": "2025-10-06T03:32:20.462964Z"
    }
   },
   "source": [
    "class Linear(bp.BrainPyObject):\n",
    "    def __init__(self, n_in, n_out):\n",
    "        super(Linear, self).__init__()\n",
    "        self.w = bm.Variable(bm.random.random((n_in, n_out)))\n",
    "        self.b = bm.Variable(bm.zeros(n_out))\n",
    "    \n",
    "    def update(self, x):\n",
    "        return x @ self.w + self.b"
   ],
   "outputs": [],
   "execution_count": 34
  },
  {
   "cell_type": "code",
   "id": "0e5eca39",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.485128Z",
     "start_time": "2025-10-06T03:32:20.480741Z"
    }
   },
   "source": [
    "x = bm.zeros(10)  # the input data\n",
    "l = Linear(10, 3)  # the class we need"
   ],
   "outputs": [],
   "execution_count": 35
  },
  {
   "cell_type": "markdown",
   "id": "f3af6f71",
   "metadata": {},
   "source": [
    "First, we mark \"w\" and \"b\" as dynamically changed variables. Changing \"w\" or \"b\" will change the final results. "
   ]
  },
  {
   "cell_type": "code",
   "id": "4cca2e8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.521715Z",
     "start_time": "2025-10-06T03:32:20.492138Z"
    }
   },
   "source": [
    "update1 = bm.jit(\n",
    "    l.update,\n",
    ")  \n",
    "\n",
    "update1(x)  # x is 0., b is 0., therefore y is 0."
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0., 0., 0.], dtype=float32)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 36
  },
  {
   "cell_type": "code",
   "id": "c4c9c2f2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.554412Z",
     "start_time": "2025-10-06T03:32:20.548111Z"
    }
   },
   "source": [
    "l.b[:] = 1.  # change b to 1, we expect y will be 1 too\n",
    "\n",
    "update1(x)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1., 1., 1.], dtype=float32)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 37
  },
  {
   "cell_type": "markdown",
   "id": "2572c4d8",
   "metadata": {},
   "source": [
    "This time, we only mark \"w\" as a dynamically changed variable. We will find that no matter how \"b\" is modified, the results will not change. "
   ]
  },
  {
   "cell_type": "code",
   "id": "c7bdf120",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.601024Z",
     "start_time": "2025-10-06T03:32:20.573280Z"
    }
   },
   "source": [
    "update2 = bm.jit(\n",
    "    l.update,\n",
    ")\n",
    "\n",
    "update2(x)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1., 1., 1.], dtype=float32)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 38
  },
  {
   "cell_type": "code",
   "id": "446ea19c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.636349Z",
     "start_time": "2025-10-06T03:32:20.631031Z"
    }
   },
   "source": [
    "l.b[:] = 2.  # change b to 2\n",
    "\n",
    "update2(x)  # while y will not be 2"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([2., 2., 2.], dtype=float32)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 39
  },
  {
   "cell_type": "markdown",
   "id": "dfb25ba8",
   "metadata": {},
   "source": [
    "#### Example 2: JIT compiled functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b424f80",
   "metadata": {},
   "source": [
    "Now, we change the above \"Linear\" object to a function. "
   ]
  },
  {
   "cell_type": "code",
   "id": "675ce89d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.659674Z",
     "start_time": "2025-10-06T03:32:20.654711Z"
    }
   },
   "source": [
    "n_in = 10;  n_out = 3\n",
    "\n",
    "w = bm.Variable(bm.random.random((n_in, n_out)))\n",
    "b = bm.Variable(bm.zeros(n_out))\n",
    "\n",
    "def update(x):\n",
    "    return x @ w + b"
   ],
   "outputs": [],
   "execution_count": 40
  },
  {
   "cell_type": "markdown",
   "id": "1a9ffb7e",
   "metadata": {},
   "source": [
    "If we do not provide ``dyn_vars``, \"w\" and \"b\" will be compiled as constant values. "
   ]
  },
  {
   "cell_type": "code",
   "id": "a5e3c1c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.698840Z",
     "start_time": "2025-10-06T03:32:20.671861Z"
    }
   },
   "source": [
    "update1 = bm.jit(update)\n",
    "update1(x)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0., 0., 0.], dtype=float32)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 41
  },
  {
   "cell_type": "code",
   "id": "922fd101",
   "metadata": {
    "scrolled": true,
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.721495Z",
     "start_time": "2025-10-06T03:32:20.714889Z"
    }
   },
   "source": [
    "b[:] = 1.  # modify the value of 'b' will not \n",
    "           # change the result, because in the \n",
    "           # jitted function, 'b' is already \n",
    "           # a constant\n",
    "update1(x)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1., 1., 1.], dtype=float32)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 42
  },
  {
   "cell_type": "markdown",
   "id": "2dbbd220",
   "metadata": {},
   "source": [
    "Providing \"w\" and \"b\" as ``dyn_vars`` will make them dynamically changed again. "
   ]
  },
  {
   "cell_type": "code",
   "id": "c301f14b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.763550Z",
     "start_time": "2025-10-06T03:32:20.734922Z"
    }
   },
   "source": [
    "update2 = bm.jit(update)\n",
    "update2(x)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1., 1., 1.], dtype=float32)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 43
  },
  {
   "cell_type": "code",
   "id": "165bb3b2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.787332Z",
     "start_time": "2025-10-06T03:32:20.781029Z"
    }
   },
   "source": [
    "b[:] = 2.  # change b to 2, while y will not be 2\n",
    "update2(x)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([2., 2., 2.], dtype=float32)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 44
  },
  {
   "cell_type": "markdown",
   "id": "4f2c54d3",
   "metadata": {},
   "source": [
    "#### Example 3: JIT compiled neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "654a0425",
   "metadata": {},
   "source": [
    "Now, let's use SGD to train a neural network with JIT acceleration. Here we use the autograd function ``brainpy.math.grad()``, which will be discussed in detail in [the next section](./differentiation.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "id": "4b89b7af",
   "metadata": {
    "lines_to_next_cell": 2,
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.823230Z",
     "start_time": "2025-10-06T03:32:20.801953Z"
    }
   },
   "source": [
    "class LinearNet(bp.BrainPyObject):\n",
    "    def __init__(self, n_in, n_out):\n",
    "        super(LinearNet, self).__init__()\n",
    "\n",
    "        # weights\n",
    "        self.w = bm.TrainVar(bm.random.random((n_in, n_out)))\n",
    "        self.b = bm.TrainVar(bm.zeros(n_out))\n",
    "        self.r = bm.TrainVar(bm.random.random((n_out, 1)))\n",
    "    \n",
    "    def update(self, x):\n",
    "        h = x @ self.w + self.b\n",
    "        return h @ self.r\n",
    "    \n",
    "    def loss(self, x, y):\n",
    "        predict = self.update(x)\n",
    "        return bm.mean((predict - y) ** 2)\n",
    "\n",
    "\n",
    "ln = LinearNet(100, 200)\n",
    "\n",
    "# provide the variables want to update\n",
    "opt = bp.optim.SGD(lr=1e-6, train_vars=ln.vars())\n",
    "\n",
    "# provide the variables require graidents\n",
    "f_grad = bm.grad(ln.loss, grad_vars=ln.vars(), return_value=True)  \n",
    "\n",
    "\n",
    "def train(X, Y):\n",
    "    grads, loss = f_grad(X, Y)\n",
    "    opt.update(grads)\n",
    "    return loss\n",
    "\n",
    "# JIT the train function \n",
    "train_jit = bm.jit(train)"
   ],
   "outputs": [],
   "execution_count": 45
  },
  {
   "cell_type": "code",
   "id": "8ae01dee",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:20.997081Z",
     "start_time": "2025-10-06T03:32:20.829639Z"
    }
   },
   "source": [
    "xs = bm.random.random((1000, 100))\n",
    "ys = bm.random.random((1000, 1))\n",
    "\n",
    "for i in range(30):\n",
    "    loss  = train_jit(xs, ys)\n",
    "    print(f'Train {i}, loss = {loss:.2f}')"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train 0, loss = 6155007.00\n",
      "Train 1, loss = 3455055.50\n",
      "Train 2, loss = 1950462.62\n",
      "Train 3, loss = 1104584.25\n",
      "Train 4, loss = 626670.44\n",
      "Train 5, loss = 355895.41\n",
      "Train 6, loss = 202236.78\n",
      "Train 7, loss = 114960.27\n",
      "Train 8, loss = 65362.71\n",
      "Train 9, loss = 37169.18\n",
      "Train 10, loss = 21140.02\n",
      "Train 11, loss = 12025.93\n",
      "Train 12, loss = 6843.44\n",
      "Train 13, loss = 3896.46\n",
      "Train 14, loss = 2220.66\n",
      "Train 15, loss = 1267.70\n",
      "Train 16, loss = 725.79\n",
      "Train 17, loss = 417.62\n",
      "Train 18, loss = 242.38\n",
      "Train 19, loss = 142.73\n",
      "Train 20, loss = 86.06\n",
      "Train 21, loss = 53.83\n",
      "Train 22, loss = 35.51\n",
      "Train 23, loss = 25.09\n",
      "Train 24, loss = 19.16\n",
      "Train 25, loss = 15.79\n",
      "Train 26, loss = 13.87\n",
      "Train 27, loss = 12.78\n",
      "Train 28, loss = 12.16\n",
      "Train 29, loss = 11.81\n"
     ]
    }
   ],
   "execution_count": 46
  },
  {
   "cell_type": "markdown",
   "id": "967345db",
   "metadata": {},
   "source": [
    "### RandomState"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26106295",
   "metadata": {},
   "source": [
    "We have talked about RandomState in the [Variables](./variables.ipynb) section. RandomeState is also a Variable. Therefore, if the default RandomState (``brainpy.math.random.DEFAULT``) is used in your function, you should mark it as one of the ``dyn_vars`` in the function. Otherwise, they will be treated as constants and the jitted function will always return the same value. "
   ]
  },
  {
   "cell_type": "code",
   "id": "fe1a5925",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.017257Z",
     "start_time": "2025-10-06T03:32:21.012850Z"
    }
   },
   "source": [
    "def function():\n",
    "    return bm.random.normal(0, 1, size=(10,))"
   ],
   "outputs": [],
   "execution_count": 47
  },
  {
   "cell_type": "code",
   "id": "93c3d479",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.150943Z",
     "start_time": "2025-10-06T03:32:21.032172Z"
    }
   },
   "source": [
    "f1 = bm.jit(function)\n",
    "\n",
    "f1() == f1()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([False, False, False, False, False, False, False, False, False,\n",
       "       False], dtype=bool)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 48
  },
  {
   "cell_type": "markdown",
   "id": "95276b4b",
   "metadata": {},
   "source": [
    "The correct way to make JIT for this function is:"
   ]
  },
  {
   "cell_type": "code",
   "id": "5dfba12e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.282527Z",
     "start_time": "2025-10-06T03:32:21.167508Z"
    }
   },
   "source": [
    "bm.random.seed(1234)\n",
    "\n",
    "f2 = bm.jit(function)\n",
    "\n",
    "f2() == f2()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([False, False, False, False, False, False, False, False, False,\n",
       "       False], dtype=bool)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 49
  },
  {
   "cell_type": "markdown",
   "id": "54f23e60",
   "metadata": {},
   "source": [
    "### Static arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf607131",
   "metadata": {},
   "source": [
    "Static arguments are treated as static/constant in the jitted function. \n",
    "\n",
    "Two things must be marked as static: numerical arguments used in the conditional syntax (bool values or resulting in bool values) and strings. Otherwise, an error will raise. "
   ]
  },
  {
   "cell_type": "code",
   "id": "c624ede7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.302735Z",
     "start_time": "2025-10-06T03:32:21.297975Z"
    }
   },
   "source": [
    "@bm.jit\n",
    "def f(x):\n",
    "  if x < 3:  # this will cause error\n",
    "    return 3. * x ** 2\n",
    "  else:\n",
    "    return -4 * x"
   ],
   "outputs": [],
   "execution_count": 50
  },
  {
   "cell_type": "code",
   "id": "43d03199",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.323303Z",
     "start_time": "2025-10-06T03:32:21.308564Z"
    }
   },
   "source": [
    "try:\n",
    "    f(1.)\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'jax.errors.TracerBoolConversionError'> Attempted boolean conversion of traced array with shape bool[].\n",
      "See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n"
     ]
    }
   ],
   "execution_count": 51
  },
  {
   "cell_type": "markdown",
   "id": "aa080dcc",
   "metadata": {},
   "source": [
    "Simply speaking, arguments resulting in boolean values must be declared as static arguments. In ``brainpy.math.jit()`` function, we can set the names of static arguments. "
   ]
  },
  {
   "cell_type": "code",
   "id": "3005cf57",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.332809Z",
     "start_time": "2025-10-06T03:32:21.327878Z"
    }
   },
   "source": [
    "def f(x):\n",
    "  if x < 3:  # this will cause error\n",
    "    return 3. * x ** 2\n",
    "  else:\n",
    "    return -4 * x\n",
    "\n",
    "f_jit = bm.jit(f, static_argnames=('x',))"
   ],
   "outputs": [],
   "execution_count": 52
  },
  {
   "cell_type": "code",
   "id": "41349cb1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:21.370385Z",
     "start_time": "2025-10-06T03:32:21.354879Z"
    }
   },
   "source": [
    "f_jit(x=1.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(3., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 53
  },
  {
   "cell_type": "markdown",
   "id": "86485a58",
   "metadata": {},
   "source": [
    "However, it's worth noting that calling the jitted function with different values for these static arguments will trigger recompilation. Therefore, declaring static arguments may be suitable to the following situations:\n",
    "\n",
    "1. Boolean arguments.\n",
    "2. Arguments that only have several possible values. \n",
    "\n",
    "If the argument value change significantly, you'd better not declare it as static. \n",
    "\n",
    "For more information, please refer to the [jax.jit](https://jax.readthedocs.io/en/latest/jax.html?highlight=jit#jax.jit) API."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "279.273px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
